import cv2
import torch
from PIL import Image
import numpy as np

from ultralytics import YOLO
from ultralytics.models.yolo.detect import DetectionPredictor
from ultralytics.engine.results import Results
from ultralytics.utils import ops
import matplotlib.patches as patches

from matplotlib import pyplot as plt

class HookDetectionPredictor(DetectionPredictor):
    """
    A class extending the BasePredictor class for prediction based on a detection model.

    Example:
        ```python
        from ultralytics.utils import ASSETS
        from ultralytics.models.yolo.detect import DetectionPredictor

        args = dict(model="yolo11n.pt", source=ASSETS)
        predictor = DetectionPredictor(overrides=args)
        predictor.predict_cli()
        ```
    """

    def postprocess(self, preds, img, orig_imgs):
        """Post-processes predictions and returns a list of Results objects."""
        self.pred = preds[0][0].transpose(0,1)
        print(self.pred.shape)
        self.pred[:, :4] = ops.scale_boxes(img.shape[2:], self.pred[:, :4], orig_imgs[0].shape)
        # print(pred.shape)
        # print(pred[0,:,1])

        for x in preds[1]:
            print(x.shape)

        preds = ops.non_max_suppression(
            preds,
            self.args.conf,
            self.args.iou,
            agnostic=self.args.agnostic_nms,
            max_det=self.args.max_det,
            classes=self.args.classes,
        )

        if not isinstance(orig_imgs, list):  # input images are a torch.Tensor, not a list
            orig_imgs = ops.convert_torch2numpy_batch(orig_imgs)

        results = []
        for pred, orig_img, img_path in zip(preds, orig_imgs, self.batch[0]):
            pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape)
            results.append(Results(orig_img, path=img_path, names=self.model.names, boxes=pred))
        return results

def draw_points(image, preds):
    fig, ax = plt.subplots(1)
    ax.imshow(image)

    alpha_map = np.zeros((image.size[1], image.size[0]), dtype=float)
    count_map = np.zeros((image.size[1], image.size[0]), dtype=float)

    for x, y, w, h, alpha in preds:
        alpha_map[int(y):int(h), int(x):int(w)] += alpha.item()
        count_map[int(y):int(h), int(x):int(w)] += 1
    #alpha_map /= alpha_map.max()
    alpha_map /= count_map
    print(alpha_map[390, 80], alpha_map[110, 720])
    # print(alpha_map[50, 200], alpha_map[665, 780])
    #print(alpha_map[410, 700])

    ax.imshow(np.full((image.size[1], image.size[0]), 255, dtype=np.uint8), alpha=alpha_map)

    #ax.scatter(preds[:,0], preds[:,1], alpha=torch.clamp(preds[:,4]*0.1, min=0.01))

    # for x, y, w, h, alpha in preds:
    #     rect = patches.Rectangle((x, y), w-x, h-y,
    #                              linewidth=1,
    #                              edgecolor=None,  # 边界框颜色
    #                              facecolor='g',  # 填充颜色
    #                              alpha=max(alpha.item()*0.02, 0.002),
    #                              fill=True)  # 启用填充
    #     ax.add_patch(rect)

    plt.show()


model = YOLO("ckpts/y1_735k_bs512_nbs512_s_yv11.pt")

# from PIL
im1 = Image.open("imgs/t1.jpg")
results = model.predict(source=im1, save=True, predictor=HookDetectionPredictor)  # save plotted images
#print(results[0].boxes)


draw_points(im1, model.predictor.pred)